#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-01-20

"""
【节点名称】：
    TaskManagerNode
【依赖项安装】：
    cd <path-to-yolo11-spire-mod>
    pip install -e .
【订阅类型】：
    无
【发布类型】：
    无
【备注】：
    无
"""

import threading
import time
import random
import os
import json
import math
from collections import defaultdict
import cv2
from typing import Union
from queue import Queue
from copy import deepcopy, copy
from pathlib import Path
import numpy as np
from spirems import Subscriber, Publisher, cvimg2sms, def_msg, QoS, Logger, Service
from spirems.nodes.BaseNode import BaseNode
import uuid
import torch
import base64
import glob
import subprocess
import threading
import yaml


stdout_log = {}
stderr_log = {}


def read_stream(stream, prefix, tid):
    """读取流并打印内容"""
    for line in stream:
        print(f"{prefix}: {line.strip()}")
        if "STDOUT" == prefix:
            if tid not in stdout_log:
                stdout_log[tid] = []
            stdout_log[tid].append(line.strip())
            if len(stdout_log[tid]) > 20:
                stdout_log[tid].pop(0)
        if "STDERR" == prefix:
            if tid not in stderr_log:
                stderr_log[tid] = []
            stderr_log[tid].append(line.strip())
            if len(stderr_log[tid]) > 20:
                stderr_log[tid].pop(0)


def run_command_with_realtime_output(command, tid):
    try:
        # 使用 subprocess.Popen 启动进程
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,  # 捕获标准输出
            stderr=subprocess.PIPE,  # 捕获标准错误
            shell=True,              # 使用 shell 执行命令
            text=True,               # 以文本模式处理输出
            encoding='iso-8859-1'
        )
        
        # 创建线程分别读取标准输出和标准错误
        stdout_thread = threading.Thread(target=read_stream, args=(process.stdout, "STDOUT", tid))
        stderr_thread = threading.Thread(target=read_stream, args=(process.stderr, "STDERR", tid))
        
        # 启动线程
        stdout_thread.start()
        stderr_thread.start()
        
        # 等待线程完成
        stdout_thread.join()
        stderr_thread.join()
        
        # 等待进程完成
        process.wait()
        
        # 获取返回码
        return_code = process.returncode
        
        # 返回返回码
        return return_code
    
    except Exception as e:
        # 捕获并返回异常
        print(f"Exception: {e}")
        return -1


class TaskManagerNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = False
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            sms_logger=True
        )
        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.sub = Subscriber("/YOLO11TrainManager/request", "spirecv_msgs::Task", self.service_callback)
        self.pub = Publisher("/YOLO11TrainManager/response", "spirecv_msgs::TaskResult")

        self.training = False
        self.before_queue = Queue()
        self.after_queue = Queue()
        self.task_id = 0
        self.task = {}
        self.start()

    def release(self):
        BaseNode.release(self)
        self.sub.kill()
        self.pub.kill()
    
    def task_list(self):
        plst  = []
        flst  = []
        cp_bqueue = list(self.before_queue.queue)
        cp_aqueue = list(self.after_queue.queue)
        cp_task = copy(self.task)
        plst.append(cp_task)
        for item in cp_bqueue:
            plst.append(item)
        for item in cp_aqueue:
            flst.append(item)
        return plst, flst

    def service_callback(self, msg):
        self.logger.info("RECV: " + msg['task'])

        res = def_msg("spirecv_msgs::TaskResult")
        res['task'] = msg['task']
        if "query" == msg['task']:
            plst, flst = self.task_list()
            res['data']['train_queue'] = plst
            res['data']['finish_queue'] = flst
        elif "train" == msg['task']:
            cmd = msg['data']
            if cmd.startswith('yolo detect train'):
                cmd = cmd[17:].strip()
            params = cmd.split(' ')
            sv_dir_train, sv_dir_val = None, None
            cmd_cp = ''
            for param in params:
                if len(param.split('=')) != 2:
                    continue
                key, val = param.split('=')
                if 'sv_dir_train' == key:
                    sv_dir_train = val
                elif 'sv_dir_val' == key:
                    sv_dir_val = val
                else:
                    cmd_cp += ' ' + param
            if sv_dir_train is not None:
                if sv_dir_val is None:
                    sv_dir_val = sv_dir_train
                yaml_data = {'train_sv': sv_dir_train, 'val_sv': sv_dir_val}
                yaml_fn = str(uuid.uuid4()).replace('-', '_') + '.yaml'
                with open(yaml_fn, 'w') as f:
                    yaml.dump(yaml_data, f, default_flow_style=False)
                cmd_cp += ' data=' + yaml_fn
            cmd = 'yolo detect train ' + cmd
            cmd_cp = 'yolo detect train' + cmd_cp
            # print("cmd", cmd)
            # print("cmd_cp", cmd_cp)
            if self.task_id >= 1000:
                self.task_id = 0
            self.task_id += 1
            self.before_queue.put({"tid": self.task_id, "cmd": cmd, "cmd_cp": cmd_cp})
        elif "kill" == msg['task']:
            try:
                _tid = int(msg['data'])
                print("kill", _tid)
                queue_keep = []
                while not self.before_queue.empty():
                    item = self.before_queue.get()
                    if item['tid'] != _tid:
                        queue_keep.append(item)
                for item in queue_keep:
                    self.before_queue.put(item)
            except:
                pass
        elif "log" == msg['task']:
            try:
                _tid = int(msg['data'])
                print("log", _tid)
                if _tid in stdout_log:
                    res['data']['stdout'] = stdout_log[_tid]
                if _tid in stderr_log:
                    res['data']['stderr'] = stderr_log[_tid]
            except:
                pass

        self.pub.publish(res)

    def run(self):
        while self.is_running():
            self.training = False
            time.sleep(1)

            self.task = {}
            if self.before_queue.qsize() > 0:
                self.task = self.before_queue.get()
                cmd = self.task["cmd_cp"]
                self.training = True
                run_command_with_realtime_output(cmd, self.task["tid"])

            if len(self.task) > 0:
                self.after_queue.put(self.task)
                while self.after_queue.qsize() > 10:
                    self.after_queue.get()

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        default=None,
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='yolo11',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args = parser.parse_args()
    print("config:", args.config)
    print("job-name:", args.job_name)

    node = TaskManagerNode(
        job_name=args.job_name,
        param_dict_or_file=args.config,
        ip=args.ip,
        port=args.port
    )
    node.join()
    print('done!')


if __name__ == '__main__':
    main()
